from ltp import LTP
import wenet
from keybert import KeyBERT


keywords_model = KeyBERT('hfl/chinese-bert-wwm-ext')   #'bert-base-chinese'
ltp = LTP(pretrained_model_name_or_path="./resource/ltp/base/", map_location='cpu')

asr_model = wenet.load_model('chinese')

#HF_ENDPOINT=https://hf-mirror.com TRANSFORMERS_CACHE=./resource/tmp/ python task_clay.py
#HF_ENDPOINT=https://hf-mirror.com TRANSFORMERS_CACHE=/Users/on/programs/mate_pretrain/resource/tmp python app.py

def proc_pipline(wav_path):
    result = asr_model.transcribe(wav_path)
    wav_asr_text = result['text']

    text = '推动实施智能教育的措施不能仅从高等学校人才培养和人工智能发展的必要性角度思考'
    text = '证监会拟决定对许家印采取终身证券市场禁入措施，对恒大地产处以 41.75 亿元罚款，将带来哪些影响？'
    text = wav_asr_text
    #text = '我讨厌孤独'
    words = ltp.pipeline(text, tasks=['cws'])['cws']
    keywords = keywords_model.extract_keywords(' '.join(words), keyphrase_ngram_range=(1,1),  top_n=5)
    print(f'asr result:{text}')
    print(f'key words:{keywords}')
    
    return keywords

if __name__ == '__main__':
    proc_pipline(wav_path = '/Users/on/Documents/wav/output2.wav')